// Public domain AES implementation using ARM Cryptography Extensions
// 适用于aarch64架构
// 支持AES128和AES256加密

#include <assert.h>
#include <string.h>
#include <arm_neon.h>
#include <arm_acle.h>

#include "aes_arm.h"

// AES轮常量
static const uint8_t rcon[11] = {
    0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36
};

// 轮密钥扩展函数 - 用于AES128
static uint8x16_t aes128_key_expand(uint8x16_t key, int round) {
    // 旋转字节
    uint8x16_t rotated = vextq_u8(key, key, 1);
    // 替换字节（使用AES S盒）
    uint8x16_t subbed = vaeseq_u8(rotated, vdupq_n_u8(0));
    subbed = vaesmcq_u8(subbed);
    // 与轮常量异或
    uint8x16_t rcon_vec = vdupq_n_u8(0);
    rcon_vec = vsetq_lane_u8(rcon[round], rcon_vec, 0);
    // 完成扩展
    return veorq_u8(key, veorq_u8(subbed, rcon_vec));
}

// 轮密钥扩展函数 - 用于AES256
static uint8x16_t aes256_key_expand(uint8x16_t key, uint8x16_t prev_key, int round) {
    if (round % 2 == 0) {
        // 旋转字节
        uint8x16_t rotated = vextq_u8(key, key, 1);
        // 替换字节
        uint8x16_t subbed = vaeseq_u8(rotated, vdupq_n_u8(0));
        subbed = vaesmcq_u8(subbed);
        // 与轮常量异或
        uint8x16_t rcon_vec = vdupq_n_u8(0);
        rcon_vec = vsetq_lane_u8(rcon[round/2], rcon_vec, 0);
        return veorq_u8(prev_key, veorq_u8(subbed, rcon_vec));
    } else {
        // 替换字节但不旋转
        uint8x16_t subbed = vaeseq_u8(key, vdupq_n_u8(0));
        subbed = vaesmcq_u8(subbed);
        return veorq_u8(prev_key, subbed);
    }
}

void AES128_load_schedule(const uint8_t *key, uint8_t *_schedule) {
    uint8x16_t *schedule = (uint8x16_t *)_schedule;
    
    // 加载初始密钥
    schedule[0] = vld1q_u8(key);
    
    // 扩展密钥
    for (int i = 1; i <= 10; i++) {
        schedule[i] = aes128_key_expand(schedule[i-1], i);
    }
}

void AES256_load_schedule(const uint8_t *key, uint8_t *_schedule) {
    uint8x16_t *schedule = (uint8x16_t *)_schedule;
    
    // 加载初始密钥（256位密钥分为两个128位块）
    schedule[0] = vld1q_u8(key);
    schedule[1] = vld1q_u8(key + 16);
    
    // 扩展密钥
    for (int i = 2; i <= 14; i++) {
        schedule[i] = aes256_key_expand(schedule[i-1], schedule[i-2], i-1);
    }
}

static inline void aes128_enc(const uint8_t *plaintext,
                             const uint8_t *_schedule, uint8_t *ciphertext) {
    const uint8x16_t *schedule = (const uint8x16_t *)_schedule;
    
    // 加载明文
    uint8x16_t state = vld1q_u8(plaintext);
    
    // 初始轮：与第一轮密钥异或
    state = veorq_u8(state, schedule[0]);
    
    // 主加密轮
    for (int i = 1; i < 10; i++) {
        state = vaeseq_u8(state, schedule[i]);
        state = vaesmcq_u8(state);
    }
    
    // 最后一轮（没有混合列操作）
    state = vaeseq_u8(state, schedule[10]);
    
    // 存储密文
    vst1q_u8(ciphertext, state);
}

static inline void aes256_enc(const uint8_t *plaintext,
                             const uint8_t *_schedule, uint8_t *ciphertext) {
    const uint8x16_t *schedule = (const uint8x16_t *)_schedule;
    
    // 加载明文
    uint8x16_t state = vld1q_u8(plaintext);
    
    // 初始轮：与第一轮密钥异或
    state = veorq_u8(state, schedule[0]);
    
    // 主加密轮
    for (int i = 1; i < 14; i++) {
        state = vaeseq_u8(state, schedule[i]);
        state = vaesmcq_u8(state);
    }
    
    // 最后一轮（没有混合列操作）
    state = vaeseq_u8(state, schedule[14]);
    
    // 存储密文
    vst1q_u8(ciphertext, state);
}

void AES128_CTR_enc_sch(const uint8_t *plaintext, const size_t plaintext_len,
                       const uint8_t *schedule, uint8_t *ciphertext) {
    assert(plaintext_len % 16 == 0);
    for (size_t block = 0; block < plaintext_len / 16; block++) {
        aes128_enc(plaintext + (16 * block), schedule, ciphertext + (16 * block));
    }
}

void AES256_CTR_enc_sch(const uint8_t *plaintext, const size_t plaintext_len,
                       const uint8_t *schedule, uint8_t *ciphertext) {
    assert(plaintext_len % 16 == 0);
    for (size_t block = 0; block < plaintext_len / 16; block++) {
        aes256_enc(plaintext + (16 * block), schedule, ciphertext + (16 * block));
    }
}

void AES128_free_schedule(uint8_t *schedule) {
    memset(schedule, 0, 16 * 11);
}

void AES256_free_schedule(uint8_t *schedule) {
    memset(schedule, 0, 16 * 15);
}
